#include <queue>
#include "session.h"
#include "../graph/graph.h"
#include "../graph/variable.h"

using std::queue;

Matrix Session::run(const NodeWrapper & node, unordered_map<Node*, Matrix> feed_map){
    Node* root = node.getNode();
    Graph* graph = Graph::getGraph();
    unordered_map<Node*, unsigned int> node2inNum = graph -> getInNumOfNodesMap(root);
    unordered_map<Node*, unsigned int> node2outNum = graph -> getOutNumOfNodesMap(root);
    // TODO 1. 初始化context里的var和placeholder
    // 2. 当一个节点的out邻居们都计算完毕，删除context中对应node键值对，以节省空间
    Context* context = new Context(node2inNum, node2outNum, feed_map);
    compute(context);
    Matrix ret_mat;
    if(root -> isOptimizer())
        ret_mat = Matrix(0, 0);
    else
        ret_mat = context -> getFlowedData(root);
    delete context;
    return ret_mat;
}

void Session::compute(Context* context){
    unordered_map<Node*, unsigned int> in_map = context -> getInNumOfNodeMap();
    unordered_map<Node*, unsigned int> out_map = context -> getOutNumOfNodeMap();
    queue<Node*> execute_queue;
    // use the nodes with input degree equaling to zero to initialize execute_queue
    for(auto iter = in_map.begin(); iter != in_map.end(); iter++)
        if(iter -> second == 0)
            execute_queue.push(iter -> first);
    while(!execute_queue.empty()) {
        auto top = execute_queue.front();
        execute_queue.pop();
        top -> compute(context);
        auto outNeighbors = top -> getNodeOut();
        auto inNeighbors = top -> getNodeIn();
        for(int i = 0; i < outNeighbors -> size(); i++){
            auto neighbor = (*outNeighbors)[i];
            if(in_map.find(neighbor) != in_map.end() && --in_map[neighbor] == 0)
                execute_queue.push(neighbor);
        }
        for(int i = 0; i < inNeighbors -> size(); i++){
            auto neighbor = (*inNeighbors)[i];
            if(out_map.find(neighbor) != out_map.end() && --out_map[neighbor] == 0)
                context -> eraseFlowedData(neighbor);
        }
    }
}
